package com.wjf.wjfaiagent.chatmemory;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.collection.ListUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.wjf.wjfaiagent.converter.MessageConverter;
import com.wjf.wjfaiagent.entity.ChatMessage;
import com.wjf.wjfaiagent.service.IChatMessageService;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

@RequiredArgsConstructor
//@Component
public class DatabaseBasedfChatMemory implements ChatMemory {
//    @Autowired
    private final IChatMessageService messageService;
    @Override
    public void add(String conversationId, List<Message> messages) {
        List<ChatMessage> insertBatch = messages.stream().map(message -> MessageConverter.toChatMessage(conversationId, message)).collect(Collectors.toList());
        messageService.saveBatch(insertBatch);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        LambdaQueryWrapper<ChatMessage> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(ChatMessage::getConversationId, conversationId)
                .orderByDesc(ChatMessage::getCreateTime)
                .last(lastN > 0, "limit " + lastN);
        List<ChatMessage> chatMessages = messageService.list(wrapper);
        if (!CollectionUtil.isEmpty(chatMessages))
            Collections.reverse(chatMessages);
        List<Message> messages = chatMessages.stream().map(MessageConverter::toMessage).collect(Collectors.toList());
        return messages;
    }

    @Override
    public void clear(String conversationId) {
        LambdaQueryWrapper<ChatMessage> queryWrapper = new LambdaQueryWrapper<>();
        queryWrapper.eq(ChatMessage::getConversationId, conversationId);
        messageService.remove(queryWrapper);
    }
}
